import torch
import torch.nn as nn
import torch.distributed as dist
import os
import json
import shutil

def to_python_float(t):
    if hasattr(t, 'item'):
        return t.item()
    else:
        return t[0]


def reduce_tensor(tensor):
    tensor = tensor.clone()
    dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
    tensor /=dist.get_world_size()

    return tensor

class AverageMeter(object):
    """Computes and stores the average and current value"""

    def __init__(self, reduce_=None):
        self.reset()
        self.reduce = reduce_

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0


    def update(self, value, num=1):
        self.val = value
        self.sum += self.val * num
        self.count += num
        self.avg = self.sum / self.count

    def get_metric(self):
        self.avg = self.sum / self.count



def accuracy(output, target, topk=(1,)):
    """Computes the accuracy over the k top predictions for the specified values of k"""
    with torch.no_grad():
        output = output.transpose(1, 2)
        output = output.contiguous().view(-1, output.size(-1))
        target = target.view(-1)
        maxk = max(topk)
        batch_size = target.size(0)
#        print(index.shape)
#        print('top@', maxk)
        _, pred = output.topk(maxk, 1, True, True)
        pred = pred.t()
#        print('pred', pred)
        correct = pred.eq(target.view(1, -1).expand_as(pred))
        res = []
        for k in topk:
            correct_k = correct[:k].view(-1).float().sum(0, keepdim=True)
            res.append(correct_k.mul_(100.0 / batch_size))
        return res

